import os
import cv2
import glob
import time
import random
import visdom
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms
from torchvision.datasets.folder import default_loader
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt


# class NumberDataset(Dataset):
#     def __init__(self, training=True):
#         if training:
#             self.samples = list(range(1, 1001))
#         else:
#             self.samples = list(range(1001, 2000))
#
#     def __len__(self):
#         return len(self.samples)
#
#     def __getitem__(self, item):
#         return self.samples[item]
class PokemonDataset(Dataset):
    def __init__(self, root, resize, mode):
        super(PokemonDataset, self).__init__()
        self.root = root
        self.resize = resize
        # class -> label
        self.name2label = {}
        for name in sorted(os.listdir(root)):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.name2label[name] = len(self.name2label.keys())
        self.image_label = self.load_images()
        self.images = [img for img in self.image_label.keys()]
        random.shuffle(self.images)
        self.labels = [self.image_label[img] for img in self.images]
        if mode == "train":
            self.images = self.images[:int(0.7 * len(self.images))]
            self.labels = self.labels[:int(0.7 * len(self.labels))]
        elif mode == "val":
            self.images = self.images[int(0.7 * len(self.images)): int(0.9 * len(self.images))]
            self.labels = self.labels[int(0.7 * len(self.labels)): int(0.9 * len(self.labels))]
        elif mode == "test":
            self.images = self.images[int(0.9 * len(self.images)):]
            self.labels = self.labels[int(0.9 * len(self.labels)):]
        else:
            print("Please check your mode!")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, item):
        image, label = self.images[item], self.labels[item]
        transform_data = transforms.Compose([
            lambda img: Image.open(img).convert("RGB"),
            transforms.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),
            transforms.RandomRotation(15),
            transforms.CenterCrop(self.resize),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = transform_data(image)
        label = torch.from_numpy(label)
        return image, label

    def load_images(self):
        data_infos = {}
        for name in self.name2label.keys():
            image_list = sorted(os.listdir(os.path.join(self.root, name)))
            for image in image_list:
                image_path = os.path.join(self.root, name, image)
                data_infos[image_path] = np.array(self.name2label[name], dtype=np.int64)
        return data_infos


if __name__ == "__main__":
    dataset_dir = "/home/stark/algo-env/datasets/pokemon"
    # pokemon_dataset_01 = ImageFolder(dataset_dir)
    # print(pokemon_dataset_01.class_to_idx)
    pokemon_dataset = PokemonDataset(dataset_dir, 224, "train")
    pokemon_loader = DataLoader(pokemon_dataset, batch_size=8, shuffle=True)
    per_batch = next(iter(pokemon_loader))
    print(per_batch[0].size())
    # plt.figure()
    # i = 1
    for per_picture in per_batch:
        print(per_picture.size())
    #     plt.subplot(2, 4, i)
    #     i += 1
    # plt.show()
    # print(pokemon_dataset.name2label)
    # print(pokemon_dataset.image_label)
    # print(pokemon_dataset.name2label)
    # for ele in pokemon_dataset.images:
    #     print(ele, pokemon_dataset.image_label[ele])
    # print(pokemon_dataset.__len__())
    # x, y = next(iter(pokemon_dataset))
    # print("sample", x.shape, y.shape, y)
    # x = x.permute(1, 2, 0).numpy()
    # plt.imshow(x)
    # plt.show()
